import os

import soundfile
import tqdm
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

from ai.utils import utils_file

inference_pipeline_single = pipeline(
    task=Tasks.voice_activity_detection,
    model='damo/speech_fsmn_vad_zh-cn-16k-common-pytorch',
    model_revision=None,
    device='cpu',
)

inference_pipeline_scp = pipeline(
    task=Tasks.voice_activity_detection,
    model='damo/speech_fsmn_vad_zh-cn-16k-common-pytorch',
    model_revision=None,
    device='cpu',
    output_dir='./output/',
)


def get_scp_file(source_path: str = '/home/work_nfs6/lhma/wash_data/sample1000/res/muise_denoise',
                 output_scp_file: str = './output/wav.scp'):
    utils_file.makedir_for_file(output_scp_file)
    for root, dirs, files in os.walk(source_path):
        res_dic = {}
        for file in files:
            if file.endswith('.wav') or file.endswith('.WAV'):
                key = file.strip().split('_')[0] + "_" + file.strip().split('_')[1]
                wav_path = os.path.join(root, file)
                res_dic[key] = wav_path
        utils_file.write_dic_to_scp(res_dic, output_scp_file)
        break


def get_segments_for_audio(
        audio_in: str = "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/vad_example.wav"):
    """
    audio_in: url, file_path
    """
    # waveform, sample_rate = soundfile.read(audio_in)
    segments_result = inference_pipeline_single(
        audio_in=audio_in)
    return segments_result


def get_segments_for_scp_file(audio_in='./output/wav.scp'):
    inference_pipeline_scp(audio_in=audio_in)


def get_segements_for_scp_func2(wav_scp_file_path: str = './output/wav.scp',
                                target_json_file_path: str = './output/segments_task.scp'):
    dic_data = utils_file.load_dic_from_scp(wav_scp_file_path)
    utils_file.makedir_for_file(target_json_file_path)
    res_dic = {}
    num = 0
    print('开始得到得到wav.scp文件对应音频的切分段, 共有{}条数据'.format(len(dic_data)))
    for k, v in tqdm.tqdm(dic_data.items(), total=len(dic_data)):
        print(k)
        output = get_segments_for_audio(audio_in=v)
        res_dic[k] = output
        num += 1
        if num >= 10:
            break
    utils_file.write_dic_to_scp(res_dic,target_json_file_path)


if __name__ == '__main__':
    """"""
    # get_scp_file("/home/work_nfs6/lhma/wash_data/sample1000/res/muise_denoise", "./output/wav.scp")
    # get_segements_for_scp_func2()
    # dic = get_segments_for_audio()
    # print(dic['text'])

    # scp_file = "/home/work_nfs6/lhma/wash_data/sample1000/res/16kwavs.scp"
    # file_list = utils_file.load_list_file(scp_file)
    # output_scp = './output/wav.scp'
    # utils_file.makedir_for_file(output_scp)
    # res_dic = {}
    # for file in file_list:
    #     file = file.strip()
    #     if file.endswith('.wav') or file.endswith('.WAV') :
    #         filename = file.strip().split('/')[-1]
    #         key = filename.strip().split('_')[0] + "_" + filename.strip().split('_')[1]
    #         wav_path = file
    #         res_dic[key] = wav_path
    # print(res_dic)
    # utils_file.write_dic_to_scp(res_dic, output_scp)
    # print(file_list)

    get_segments_for_scp_file('./output/wav.scp')
